{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0400d899-0658-4c15-ba9d-3666d3f1e808",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "一个星期有7天。\n"
     ]
    }
   ],
   "source": [
    "# 导入所需的库和模块\n",
    "import numpy as np  # 导入NumPy库，用于处理数组和矩阵\n",
    "import torch  # 导入PyTorch库，用于深度学习\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig  # 从transformers库中导入所需的类\n",
    "\n",
    "# 指定预训练模型的名称\n",
    "model_name = \"Qwen/Qwen2-0.5B-Instruct\"\n",
    "\n",
    "# 加载指定的预训练模型，并将其移动到指定的设备\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, dtype=\"auto\")\n",
    "model.eval()  # 设置模型为推理模式，不启用梯度计算\n",
    "\n",
    "# 打印模型结构（已注释掉）\n",
    "# print(model)\n",
    "\n",
    "# 加载与模型对应的分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "# 加载模型的配置信息\n",
    "config = AutoConfig.from_pretrained(model_name)\n",
    "\n",
    "# 定义用户输入的提示\n",
    "prompt = \"一周有几天?\"\n",
    "\n",
    "# 定义消息列表，其中包括系统的角色设定和用户的输入内容\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},  # 系统角色设定为\"你是一个有帮助的助手\"\n",
    "    {\"role\": \"user\", \"content\": prompt},  # 用户输入的内容\n",
    "]\n",
    "\n",
    "# 使用分词器将消息列表转换为模型输入格式，并生成文本\n",
    "text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "# 将文本转换为模型输入的张量，并加载到指定的设备上，准备运行\n",
    "model_inputs = tokenizer([text], return_tensors=\"pt\")\n",
    "\n",
    "# 模型的输出，这里指的是模型输出的文本的 ID\n",
    "# max_new_tokens 用来控制回答的字数\n",
    "generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=10)\n",
    "\n",
    "# 对生成的ID进行切片，以去除输入ID部分，保留生成的新内容\n",
    "generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]\n",
    "\n",
    "# 使用分词器将生成的ID解码为文本，并跳过特殊标记\n",
    "response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
    "\n",
    "# 打印生成的文本\n",
    "print(response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dcce393-c11f-4db3-a688-2099b3fc8bcb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
